import sys
sys.path.append('../BNinference')

from bayesNet import *

def mleFromCompleteData(network, exampleTable):
    """Estimate probability parameters for given network using maximum 
    likelihood estimation from complete data. 
    
    Keyword arguments:
    network -- Bayesian network structure as a BN instance.
    exampleTable -- An ExampleTable containing instances of the data described 
    by the network.
        
    """
    
    def emptyCounter(nodes):
        """Return a dictionary of possible node instantiations with values 
        set to zero.
        
        Keyword arguments:
        nodes -- list of strings containing node names.
        
        """
        def instantiations(nodes):
            """Return a list of possible node instantiations.
            
            Keyword arguments:
            nodes -- list of strings containing node names.
            
            """
            if len(nodes)==1:
                return [[True],[False]]
            else:
                rest = instantiations(nodes[1:])
                return [[True]+inst for inst in rest]+\
                       [[False]+inst for inst in rest]
        return dict([(tuple(inst),0) for inst in instantiations(nodes)])
    
    # dictionary of variable indices in data 
    dataIndex = dict([(d,i) for i,d in enumerate(exampleTable.variables)])
    
#    l = [(k,v) for k,v in emptyCounter(network.nodes).iteritems()]
#    l.sort(reverse=True)
#    for k,v in l:
#        print k,v
    
    # initialize counters for all instantiations needed
    counts = {}
    for node in network.nodes:
        counts[tuple(network.probabilities[node].nodes)] = \
            emptyCounter(network.probabilities[node].nodes)
        if len(network.parents(node))>0:
            counts[tuple(network.parents(node))] = \
                emptyCounter(network.parents(node))
        
    # count the instantiations 
    for line in exampleTable.data:
        for nodes in counts.keys():
            inst = tuple([line[dataIndex[node]] for node in nodes])
            counts[nodes][inst] += 1.
    
    # compute probabilities
    for node in network.nodes:
        relNodes = network.probabilities[node].nodes
        nodeI = relNodes.index(node)
        parents = network.parents(node)
        
        if len(parents)==0:
            probs = [list(inst) + [counts[tuple(relNodes)][inst] / len(exampleTable.data)] 
                     for inst in counts[tuple(relNodes)].keys()]
            probs.sort(reverse=True)
            network.probabilities[node].probabilityTable = probs
        else:
            probs = list()
            for inst in counts[tuple(relNodes)].keys():
                try:
                    val = list(inst) + [counts[tuple(relNodes)][inst] / counts[tuple(parents)][drop(inst, nodeI)]]
                except ZeroDivisionError:
                    val = 0
                probs.append(val)
            probs.sort(reverse=True)
            network.probabilities[node].probabilityTable = probs

if __name__=="__main__":
    def getExample():
        nodes = ['Health_Aware', 'Smokes', 'Exercises']
        ptA = PT(['Health_Aware'], [])
        ptB = PT(['Health_Aware', 'Smokes'], [])
        ptC = PT(['Health_Aware', 'Exercises'], [])
        
        network = BN(nodes, {'Health_Aware':ptA, 'Smokes':ptB, 'Exercises':ptC})
        
        data = ExampleTable(nodes)
        examples = [Example(vals) for vals in
        [True, False, True], [True, False, True], [False, True, False], 
        [False, False, True], [True, False, False], [True, False, True], 
        [False, False, False], [True, False, True], [True, False, True], 
        [False, False, True], [True, False, True], [True, True, True], 
        [True, False, True], [True, True, True], [True, False, True], 
        [True, False, True]
        ]
        data.extend(examples)
        
        return network, data

    network, data = getExample()
    
    mleFromCompleteData(network, data)
    
    for node in network.nodes:
        print node
        print network.probabilities[node]